import uuid
from lib.lib import *
import json
import socket
from datetime import datetime
import asyncio
import websockets
import secrets
from lib.Event import Event

default = {
    "server_name": static["hostname"],
    "server_port": 11840,
    "broadcast_port": 12840,
    "broadcast_interval": 1,
    "server_uuid": str(uuid.uuid4()),
    "features": {},
    "auth_type": "none",
    "auth_timeout": 60,
    "password": "afedium",
    "captcha_length": 6,
}

config = Config
verification_codes = {}

Info = {
    "name": "WS服务器",
    "id": "ws_server",
    "dependencies": [],
    "pip_dependencies": ["websockets"],
    "linux_dependencies": []
}


class WsServerLogic:
    def __init__(self):
        global config
        static["event_handler"].register_event("ExternalIO_OUT", self.handle_external_output)
        config = Config("ws_server", default)

    async def _websocket_handler(self, websocket):
        try:
            if config.conf["auth_type"].lower() == "captcha":
                verification_code = ''.join(secrets.choice('0123456789') for _ in range(config.conf["captcha_length"]))
                verification_codes[websocket] = verification_code
                print(f"新客户端连接，验证码：{verification_code}")

                await websocket.send(chr(0x90) + "请输入验证码")

                client_response = await asyncio.wait_for(websocket.recv(), timeout=config.conf["auth_timeout"])

                if client_response[0] == chr(0x90) and client_response[1:] == verification_code:
                    del verification_codes[websocket]
                    await websocket.send(chr(0x91))
                    if static["debugging"]: print(f"客户端{websocket.remote_address}认证成功")
                else:
                    await websocket.send(chr(0x92))
                    if static["debugging"]: print(f"客户端{websocket.remote_address}认证失败")
                    await websocket.close()
                    return
            elif config.conf["auth_type"].lower() == "password":
                await websocket.send(chr(0x90) + "请输入密码")

                client_response = await asyncio.wait_for(websocket.recv(), timeout=config.conf["auth_timeout"])

                if client_response[0] == chr(0x90) and client_response[1:] == config.conf["password"]:
                    await websocket.send(chr(0x91))
                    if static["debugging"]: print(f"客户端{websocket.remote_address}认证成功")
                else:
                    await websocket.send(chr(0x92))
                    if static["debugging"]: print(f"客户端{websocket.remote_address}认证失败")
                    await websocket.close()
                    return

            while True:
                message = await websocket.recv()

                if static["debugging"]:
                    print(f"[WsServer] 收到客户端消息: {repr(message[:100])}...")

                static["event_handler"].trigger_event(
                    Event("ExternalIO_IN", message=message, client_id=websocket)
                )

        except asyncio.TimeoutError:
            if static["debugging"]: print(f"验证超时 {websocket.remote_address}")
            await websocket.send(chr(0x93))
        except websockets.exceptions.ConnectionClosedOK:
            if static["debugging"]: print(f"客户端 {websocket.remote_address} 正常断开连接。")
        except websockets.exceptions.ConnectionClosedError as e:
            if static["debugging"]: print(f"客户端 {websocket.remote_address} 异常断开连接: {e}")
        except Exception as e:
            if static["debugging"]: print(f"[WsServer] 连接处理异常: {e}")
            try:
                await websocket.send(f"服务器内部错误: {e}".encode())
            except Exception:
                pass
        finally:
            verification_codes.pop(websocket, None)

    def handle_external_output(self, event: Event):
        response_data = event.data.get("response_data")
        client_websocket = event.data.get("client_id")

        if client_websocket and response_data is not None:
            if static["debugging"]:
                print(
                    f"[WsServer] 发送响应给 {client_websocket.remote_address}: {repr(response_data[:100])}...")
            asyncio.create_task(client_websocket.send(response_data))
        else:
            if static["debugging"]: print(f"[WsServer] 事件数据不完整或 websocket 已关闭。")

    async def get_local_ip(self):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except:
            return socket.gethostbyname(socket.gethostname())

    async def broadcast_server(self):
        loop = asyncio.get_running_loop()
        transport, protocol = await loop.create_datagram_endpoint(
            lambda: BroadcastProtocol(),
            family=socket.AF_INET,
            allow_broadcast=True
        )

        local_ip = await self.get_local_ip()

        try:
            using_auth = config.conf["auth_type"].lower() != "none"
            while static["running"][Info["id"]]:
                data = json.dumps({
                    "server_name": config.conf["server_name"], "service": "afedium_server",
                    "ip": local_ip, "port": config.conf["server_port"],
                    "UUID": config.conf["server_uuid"],
                    "feature": static["features"],
                    "using_auth": using_auth,
                    "auth_timeout": config.conf["auth_timeout"],
                    "timestamp": datetime.now().isoformat(),
                }).encode()
                transport.sendto(data, ('255.255.255.255', config.conf["broadcast_port"]))
                await asyncio.sleep(config.conf["broadcast_interval"])
        finally:
            transport.close()

    async def run_server(self):
        # 启动广播任务
        broadcast_task = asyncio.create_task(self.broadcast_server())

        # 启动 WebSocket 服务器
        server = await websockets.serve(
            self._websocket_handler,
            "0.0.0.0",
            config.conf["server_port"],
            ping_interval=30,
            ping_timeout=60,
            close_timeout=10
        )
        if static["debugging"]: print(f" WebSocket server started on port {config.conf['server_port']}")

        try:
            while static["running"][Info["id"]]:
                await asyncio.sleep(1)
        except asyncio.CancelledError:
            pass
        finally:
            broadcast_task.cancel()
            await server.close()
            if static["debugging"]: print(f" WebSocket server stopped.")


class BroadcastProtocol(asyncio.DatagramProtocol):
    def connection_made(self, transport):
        self.transport = transport


ws_server_logic_instance = None


def __init__():
    global ws_server_logic_instance
    static["features"] = {}
    ws_server_logic_instance = WsServerLogic()
    if config.conf["auth_type"].lower() not in ["none", "password", "captcha"]:
        print(f"Invalid auth_type: {config.conf['auth_type']} in none, password, captcha")
    static["running"][Info["id"]] = True
    try:
        static["features"].update(config.conf["features"])
        asyncio.run(ws_server_logic_instance.run_server())
    except Exception as e:
        if static["debugging"]: print(f" ws_server asyncio loop error: {e}")
